#! /usr/bin/env python3
# -*- coding: utf-8 -*-

import socket
import asyncio
import logging
import websockets

LOG = logging.getLogger(__name__)


class WSClient():
    """
    WebSocket 客户端
    """

    def __init__(self, url, **kwargs):
        self.url = url
        self.callback = kwargs.get("callback")
        self.sleep_time = kwargs.get("sleep_time") or 5
        self.ping_timeout = kwargs.get("ping_timeout") or 5
        self.reply_timeout = kwargs.get("reply_timeout") or 10

    async def listen_forever(self):
        while True:
            LOG.debug("正在创建新连接……")
            try:
                async with websockets.connect(self.url) as ws:
                    while True:
                        try:
                            reply = await asyncio.wait_for(ws.recv(), timeout=self.reply_timeout)
                        except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
                            try:
                                pong = await ws.ping()
                                await asyncio.wait_for(pong, timeout=self.ping_timeout)
                                # await ws.ping()
                                logging.debug("Ping OK, 连接正常")
                                continue
                            except:
                                logging.warning(f"Ping 错误，{self.sleep_time} 后重试（取消请按 Ctrl+C）")
                                await asyncio.sleep(self.sleep_time)
                                break
                        logging.debug(f"> {reply}")
                        if self.callback:
                            self.callback(reply)
            except socket.gaierror:
                logging.error(f"Socket 错误，{self.sleep_time} 后重试（取消请按 Ctrl+C）")
                await asyncio.sleep(self.sleep_time)
                continue
            except ConnectionRefusedError:
                logging.error("连接被拒绝，请确认地址端口正确")
                logging.error(f"{self.sleep_time} 后重试（取消请按 Ctrl+C）")
                await asyncio.sleep(self.sleep_time)
                continue
            except Exception as e:
                print(e)
                continue


def start_ws_client(client):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(client.listen_forever())


def callback_fn(data, *args, **kwargs):
    # 消息处理回调
    print(f"收到消息：{data.decode('utf-8')}")


if __name__ == "__main__":
    ws_client = WSClient(url="ws://127.0.0.1:8080/ws", callback=callback_fn)
    start_ws_client(ws_client)
